"""工具注册表模块，管理和执行各类工具"""

import logging
import time
import uuid
from typing import Dict, Any, Optional, Callable, List, Tuple, Union
from abc import ABC, abstractmethod
from collections import defaultdict
import threading

logger = logging.getLogger('tool_registry')

class BaseTool(ABC):
    """工具基类，所有工具都需要继承这个类"""
    
    @property
    @abstractmethod
    def name(self) -> str:
        """工具名称"""
        pass
    
    @property
    @abstractmethod
    def description(self) -> str:
        """工具描述"""
        pass
    
    @property
    @abstractmethod
    def parameters(self) -> Dict[str, Any]:
        """工具参数定义，使用JSON Schema格式"""
        pass
    
    @abstractmethod
    def execute(self, **kwargs) -> Dict[str, Any]:
        """执行工具
        
        Args:
            **kwargs: 工具参数
            
        Returns:
            执行结果
        """
        pass
    
    def validate_parameters(self, **kwargs) -> Tuple[bool, str]:
        """验证工具参数
        
        Args:
            **kwargs: 工具参数
            
        Returns:
            (是否有效, 错误信息)
        """
        # 简单的验证逻辑，实际应用中可能需要更复杂的验证
        required_params = [k for k, v in self.parameters.get('properties', {}).items()
                          if v.get('required', False) or 'required' in self.parameters and k in self.parameters['required']]
        
        for param in required_params:
            if param not in kwargs:
                return False, f"缺少必需参数: {param}"
        
        # 类型验证
        for param_name, param_schema in self.parameters.get('properties', {}).items():
            if param_name in kwargs:
                param_value = kwargs[param_name]
                param_type = param_schema.get('type')
                
                if param_type and param_value is not None:
                    # 简单的类型检查
                    if param_type == 'string' and not isinstance(param_value, str):
                        return False, f"参数 {param_name} 应为字符串类型"
                    elif param_type == 'number' and not isinstance(param_value, (int, float)):
                        return False, f"参数 {param_name} 应为数字类型"
                    elif param_type == 'integer' and not isinstance(param_value, int):
                        return False, f"参数 {param_name} 应为整数类型"
                    elif param_type == 'boolean' and not isinstance(param_value, bool):
                        return False, f"参数 {param_name} 应为布尔类型"
                    elif param_type == 'array' and not isinstance(param_value, list):
                        return False, f"参数 {param_name} 应为数组类型"
                    elif param_type == 'object' and not isinstance(param_value, dict):
                        return False, f"参数 {param_name} 应为对象类型"
        
        return True, ""

class ToolRegistry:
    """工具注册表，管理所有可用的工具"""
    
    _instance = None
    _lock = threading.RLock()
    
    def __new__(cls):
        with cls._lock:
            if cls._instance is None:
                cls._instance = super(ToolRegistry, cls).__new__(cls)
                cls._instance._initialize()
        return cls._instance
    
    def _initialize(self):
        """初始化工具注册表"""
        # 存储工具实例
        self._tools = {}
        
        # 存储工具执行状态
        self._execution_status = {}
        
        # 白名单和黑名单
        self._whitelist = set()
        self._blacklist = set()
        
        # 执行统计
        self._execution_stats = defaultdict(lambda: {
            'count': 0,
            'total_time': 0,
            'success_count': 0,
            'error_count': 0
        })
        
        # 锁
        self._registry_lock = threading.RLock()
        self._execution_lock = threading.RLock()
        
        logger.info("工具注册表已初始化")
    
    def register_tool(self, tool: BaseTool) -> None:
        """注册工具
        
        Args:
            tool: 工具实例
        """
        if not isinstance(tool, BaseTool):
            raise TypeError("工具必须是BaseTool的实例")
        
        with self._registry_lock:
            self._tools[tool.name] = tool
            logger.info(f"已注册工具: {tool.name} - {tool.description}")
    
    def unregister_tool(self, tool_name: str) -> bool:
        """注销工具
        
        Args:
            tool_name: 工具名称
            
        Returns:
            是否成功注销
        """
        with self._registry_lock:
            if tool_name in self._tools:
                del self._tools[tool_name]
                logger.info(f"已注销工具: {tool_name}")
                return True
            logger.warning(f"尝试注销不存在的工具: {tool_name}")
            return False
    
    def get_tool(self, tool_name: str) -> Optional[BaseTool]:
        """获取工具实例
        
        Args:
            tool_name: 工具名称
            
        Returns:
            工具实例，如果不存在则返回None
        """
        return self._tools.get(tool_name)
    
    def list_tools(self, include_disabled: bool = False) -> List[Dict[str, Any]]:
        """列出所有工具
        
        Args:
            include_disabled: 是否包含禁用的工具
            
        Returns:
            工具列表
        """
        result = []
        with self._registry_lock:
            for tool_name, tool in self._tools.items():
                # 检查是否在黑名单中
                is_disabled = tool_name in self._blacklist
                
                if include_disabled or not is_disabled:
                    tool_info = {
                        'name': tool.name,
                        'description': tool.description,
                        'parameters': tool.parameters,
                        'enabled': not is_disabled,
                        'stats': self._execution_stats.get(tool_name, {
                            'count': 0,
                            'total_time': 0,
                            'success_count': 0,
                            'error_count': 0
                        })
                    }
                    result.append(tool_info)
        
        # 按名称排序
        result.sort(key=lambda x: x['name'])
        return result
    
    def set_whitelist(self, tools: List[str]) -> None:
        """设置工具白名单
        
        Args:
            tools: 工具名称列表
        """
        with self._registry_lock:
            self._whitelist = set(tools)
            logger.info(f"已设置工具白名单: {', '.join(tools)}")
    
    def set_blacklist(self, tools: List[str]) -> None:
        """设置工具黑名单
        
        Args:
            tools: 工具名称列表
        """
        with self._registry_lock:
            self._blacklist = set(tools)
            logger.info(f"已设置工具黑名单: {', '.join(tools)}")
    
    def is_tool_allowed(self, tool_name: str) -> bool:
        """检查工具是否允许使用
        
        Args:
            tool_name: 工具名称
            
        Returns:
            是否允许使用
        """
        # 如果白名单不为空，则只有在白名单中的工具才允许使用
        if self._whitelist and tool_name not in self._whitelist:
            return False
        
        # 如果在黑名单中，则不允许使用
        if tool_name in self._blacklist:
            return False
        
        # 检查工具是否存在
        return tool_name in self._tools
    
    def execute_tool(self, tool_name: str, **kwargs) -> Dict[str, Any]:
        """执行工具
        
        Args:
            tool_name: 工具名称
            **kwargs: 工具参数
            
        Returns:
            执行结果
        """
        # 检查工具是否允许使用
        if not self.is_tool_allowed(tool_name):
            error_msg = f"工具 {tool_name} 不允许使用或不存在"
            logger.warning(error_msg)
            return {
                'success': False,
                'error': error_msg,
                'execution_id': str(uuid.uuid4()),
                'tool_name': tool_name
            }
        
        # 获取工具实例
        tool = self._tools[tool_name]
        
        # 生成执行ID
        execution_id = str(uuid.uuid4())
        
        # 记录执行开始
        start_time = time.time()
        
        with self._execution_lock:
            self._execution_status[execution_id] = {
                'tool_name': tool_name,
                'status': 'running',
                'started_at': start_time,
                'parameters': kwargs,
                'result': None
            }
        
        try:
            # 验证参数
            is_valid, error_msg = tool.validate_parameters(**kwargs)
            if not is_valid:
                logger.warning(f"工具 {tool_name} 参数验证失败: {error_msg}")
                self._update_execution_status(execution_id, {
                    'status': 'error',
                    'error': error_msg,
                    'finished_at': time.time()
                })
                self._update_execution_stats(tool_name, False, time.time() - start_time)
                return {
                    'success': False,
                    'error': error_msg,
                    'execution_id': execution_id,
                    'tool_name': tool_name
                }
            
            # 执行工具
            logger.info(f"开始执行工具: {tool_name}, execution_id={execution_id}")
            result = tool.execute(**kwargs)
            
            # 确保结果是字典格式
            if not isinstance(result, dict):
                result = {
                    'result': result
                }
            
            # 记录执行完成
            end_time = time.time()
            execution_time = end_time - start_time
            
            self._update_execution_status(execution_id, {
                'status': 'completed',
                'result': result,
                'finished_at': end_time,
                'execution_time': execution_time
            })
            
            self._update_execution_stats(tool_name, True, execution_time)
            
            logger.info(f"工具执行完成: {tool_name}, execution_id={execution_id}, time={execution_time:.2f}s")
            
            # 返回结果，包含执行信息
            return {
                'success': True,
                'result': result,
                'execution_id': execution_id,
                'tool_name': tool_name,
                'execution_time': execution_time
            }
            
        except Exception as e:
            # 记录执行错误
            end_time = time.time()
            execution_time = end_time - start_time
            error_msg = str(e)
            
            self._update_execution_status(execution_id, {
                'status': 'error',
                'error': error_msg,
                'finished_at': end_time,
                'execution_time': execution_time
            })
            
            self._update_execution_stats(tool_name, False, execution_time)
            
            logger.error(f"工具执行失败: {tool_name}, execution_id={execution_id}, error={error_msg}")
            
            return {
                'success': False,
                'error': error_msg,
                'execution_id': execution_id,
                'tool_name': tool_name,
                'execution_time': execution_time
            }
    
    def _update_execution_status(self, execution_id: str, updates: Dict[str, Any]) -> None:
        """更新执行状态
        
        Args:
            execution_id: 执行ID
            updates: 要更新的状态字段
        """
        with self._execution_lock:
            if execution_id in self._execution_status:
                self._execution_status[execution_id].update(updates)
    
    def _update_execution_stats(self, tool_name: str, success: bool, execution_time: float) -> None:
        """更新执行统计
        
        Args:
            tool_name: 工具名称
            success: 是否成功
            execution_time: 执行时间
        """
        with self._execution_lock:
            stats = self._execution_stats[tool_name]
            stats['count'] += 1
            stats['total_time'] += execution_time
            if success:
                stats['success_count'] += 1
            else:
                stats['error_count'] += 1
    
    def get_execution_status(self, execution_id: str) -> Optional[Dict[str, Any]]:
        """获取工具执行状态
        
        Args:
            execution_id: 执行ID
            
        Returns:
            执行状态，如果不存在则返回None
        """
        with self._execution_lock:
            status = self._execution_status.get(execution_id)
            return status.copy() if status else None
    
    def list_execution_statuses(self, status: Optional[str] = None) -> Dict[str, Dict[str, Any]]:
        """列出工具执行状态
        
        Args:
            status: 可选的状态筛选
            
        Returns:
            执行ID到状态信息的映射
        """
        with self._execution_lock:
            result = {}
            for exec_id, exec_status in self._execution_status.items():
                if status is None or exec_status.get('status') == status:
                    result[exec_id] = exec_status.copy()
            return result
    
    def get_tool_stats(self, tool_name: str) -> Optional[Dict[str, Any]]:
        """获取工具执行统计
        
        Args:
            tool_name: 工具名称
            
        Returns:
            统计信息，如果工具不存在则返回None
        """
        if tool_name not in self._tools:
            return None
        
        with self._execution_lock:
            stats = self._execution_stats.get(tool_name, {
                'count': 0,
                'total_time': 0,
                'success_count': 0,
                'error_count': 0
            })
            
            # 计算平均执行时间
            if stats['count'] > 0:
                stats['avg_time'] = stats['total_time'] / stats['count']
            else:
                stats['avg_time'] = 0
            
            # 计算成功率
            if stats['count'] > 0:
                stats['success_rate'] = stats['success_count'] / stats['count']
            else:
                stats['success_rate'] = 0
            
            return stats.copy()
    
    def get_all_stats(self) -> Dict[str, Dict[str, Any]]:
        """获取所有工具的执行统计
        
        Returns:
            工具名称到统计信息的映射
        """
        result = {}
        with self._execution_lock:
            for tool_name in self._tools:
                stats = self.get_tool_stats(tool_name)
                if stats:
                    result[tool_name] = stats
        return result
    
    def cleanup_old_executions(self, max_age_seconds: int = 3600) -> int:
        """清理旧的执行记录
        
        Args:
            max_age_seconds: 最大保留时间（秒），默认为1小时
            
        Returns:
            清理的记录数量
        """
        current_time = time.time()
        to_cleanup = []
        
        with self._execution_lock:
            for exec_id, exec_status in self._execution_status.items():
                # 检查是否已完成或失败，并且已超过保留时间
                if exec_status.get('status') in ['completed', 'error']:
                    finished_at = exec_status.get('finished_at', 0)
                    if current_time - finished_at > max_age_seconds:
                        to_cleanup.append(exec_id)
            
            # 清理记录
            for exec_id in to_cleanup:
                del self._execution_status[exec_id]
            
            logger.info(f"已清理 {len(to_cleanup)} 条旧的工具执行记录")
            return len(to_cleanup)
    
    def enable_tool(self, tool_name: str) -> bool:
        """启用工具
        
        Args:
            tool_name: 工具名称
            
        Returns:
            是否成功启用
        """
        with self._registry_lock:
            if tool_name in self._tools and tool_name in self._blacklist:
                self._blacklist.remove(tool_name)
                logger.info(f"已启用工具: {tool_name}")
                return True
            return False
    
    def disable_tool(self, tool_name: str) -> bool:
        """禁用工具
        
        Args:
            tool_name: 工具名称
            
        Returns:
            是否成功禁用
        """
        with self._registry_lock:
            if tool_name in self._tools and tool_name not in self._blacklist:
                self._blacklist.add(tool_name)
                logger.info(f"已禁用工具: {tool_name}")
                return True
            return False

# 创建全局工具注册表实例
tool_registry = ToolRegistry()

__all__ = [
    'BaseTool',
    'ToolRegistry',
    'tool_registry'
]